//
//  NSObject+main.m
//  Test
//
//  Created by 张木林 on 2/18/25.
//khkkhk


//
//  Optimized main.mm
//

#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#include <time.h>
#include <QuartzCore/QuartzCore.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <CoreVideo/CoreVideo.h>

// 禁用 VSync 以减少 CPU-GPU 延迟
void disableVSync() {
    CVDisplayLinkRef displayLink;
    CVDisplayLinkCreateWithActiveCGDisplays(&displayLink);
    CVDisplayLinkStop(displayLink);
    CVDisplayLinkRelease(displayLink);
}

// 读取 Matrix Market 数据
void loadMatrixMarket(const char* filePath,
                      NSMutableArray<NSNumber*> *values,
                      NSMutableArray<NSNumber*> *rowIndices,
                      NSMutableArray<NSNumber*> *colIndices,
                      int* numRows,
                      int* numCols) {
    FILE* file = fopen(filePath, "r");
    if (!file) {
        perror("Error opening file");
        return;
    }

    char line[1024];
    *numRows = 0;
    *numCols = 0;

    while (fgets(line, sizeof(line), file)) {
        if (line[0] == '%') continue; // 忽略注释行

        int row, col;
        float value;
        if (sscanf(line, "%d %d %f", &row, &col, &value) == 3) {
            [rowIndices addObject:@(row - 1)];
            [colIndices addObject:@(col - 1)];
            [values addObject:@(value)];
            if (row > *numRows) *numRows = row;
            if (col > *numCols) *numCols = col;
        }
    }

    fclose(file);
}

// **优化的 Kernel 执行**
auto executeKernelWithTiming = ^(NSString* kernelName,
                                 NSArray<id<MTLBuffer>> *buffers,
                                 id<MTLCommandQueue> commandQueue,
                                 int numRows) {
    //id<MTLFunction> function = [commandQueue.device newDefaultLibrary].newFunctionWithName(kernelName);
    id<MTLLibrary>library = [commandQueue.device newDefaultLibrary];
    id<MTLFunction>function = [library newFunctionWithName:kernelName];
    NSError *error = nil;
    id<MTLComputePipelineState> pipeline = [commandQueue.device newComputePipelineStateWithFunction:function error:&error];
    
    if (!pipeline) {
        NSLog(@"[ERROR] Compute pipeline not found: %@", kernelName);
        return;
    }

    id<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];
    id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
    [encoder setComputePipelineState:pipeline];

    for (int i = 0; i < buffers.count; i++) {
        [encoder setBuffer:buffers[i] offset:0 atIndex:i];
    }

    // ✅ 让 threadgroupSize 自适应矩阵大小
    NSUInteger threadsPerThreadgroup = pipeline.maxTotalThreadsPerThreadgroup / 4; // 避免超载
    MTLSize threadgroupSize = MTLSizeMake(threadsPerThreadgroup, 1, 1);
    MTLSize threadgroups = MTLSizeMake((numRows + threadsPerThreadgroup - 1) / threadsPerThreadgroup, 1, 1);
    [encoder dispatchThreadgroups:threadgroups threadsPerThreadgroup:threadgroupSize];

    [encoder endEncoding];
    [commandBuffer commit];
    [commandBuffer addCompletedHandler:^(id<MTLCommandBuffer> buffer) {
        NSLog(@"%@ Execution Completed", kernelName);
    }];
};






int main() {
    @autoreleasepool {
        
        //这是一起的修改
        //禁用VSync
        disableVSync();
        // 初始化 Metal
        id<MTLDevice> device = MTLCreateSystemDefaultDevice();
        id<MTLCommandQueue> commandQueue = [device newCommandQueue];

        // ========================
        // 加载数据集
        // ========================
        NSMutableArray<NSNumber*> *values = [NSMutableArray array];
        NSMutableArray<NSNumber*> *rowIndices = [NSMutableArray array];
        NSMutableArray<NSNumber*> *colIndices = [NSMutableArray array];
        int numRows = 0, numCols = 0;

        // 替换为您的文件路径
       // const char* filePath = "/Users/zhangmulin/Downloads/kmer_V2a/kmer_V2a.mtx";
        const char* filePath = "/Users/zhangmulin/Downloads/FastLoad-main/Data/1138_bus.mtx";
        loadMatrixMarket(filePath, values, rowIndices, colIndices, &numRows, &numCols);

        // 将 NSMutableArray 转换为原始 C 数组
        int nnz = (int)values.count;
        
        size_t chunkSize = 50 * 1024 * 1024;//每块50mb
        size_t requiredSize = nnz * sizeof(float);
        size_t numChunks = (requiredSize + chunkSize - 1)/chunkSize;
        
        NSMutableArray<id<MTLBuffer>> * csrValueChunks = [NSMutableArray array];
        
        for(size_t i = 0;i<numChunks;i++)
        {
            size_t thisChunkSize = MIN(chunkSize,requiredSize - i * chunkSize);
            id<MTLBuffer>buffer = [device newBufferWithLength:thisChunkSize options:MTLResourceStorageModeShared];
            [csrValueChunks addObject:buffer];
        }
        float* c_values = (float*)malloc(nnz * sizeof(float));
        int* c_rowIndices = (int*)malloc(nnz * sizeof(int));
        int* c_colIndices = (int*)malloc(nnz * sizeof(int));

        for (int i = 0; i < nnz; i++) {
            c_values[i] = [values[i] floatValue];
            c_rowIndices[i] = [rowIndices[i] intValue];
            c_colIndices[i] = [colIndices[i] intValue];
        }

        // 转换为 CSR 格式（计算行指针）
        int* c_rowPointers = (int*)malloc((numRows + 1) * sizeof(int));
        memset(c_rowPointers, 0, (numRows + 1) * sizeof(int));
        for (int i = 0; i < nnz; i++) {
            c_rowPointers[c_rowIndices[i] + 1]++;
        }
        for (int i = 1; i <= numRows; i++) {
            c_rowPointers[i] += c_rowPointers[i - 1];
        }

        // 转换为 CSC 格式
        int* c_cscColPointers = (int*)malloc((numCols + 1) * sizeof(int));
        int* c_cscRowIndices = (int*)malloc(nnz * sizeof(int));
        float* c_cscValues = (float*)malloc(nnz * sizeof(float));
        memset(c_cscColPointers, 0, (numCols + 1) * sizeof(int));
        
        for (int i = 0; i < nnz; i++) {
            c_cscColPointers[c_colIndices[i] + 1]++;
        }
        for (int i = 1; i <= numCols; i++) {
            c_cscColPointers[i] += c_cscColPointers[i - 1];
        }
        for (int i = 0; i < nnz; i++) {
            int col = c_colIndices[i];
            int dest = c_cscColPointers[col]++;
            c_cscRowIndices[dest] = c_rowIndices[i];
            c_cscValues[dest] = c_values[i];
        }
        for (int i = numCols; i > 0; i--) {
            c_cscColPointers[i] = c_cscColPointers[i - 1];
        }
        c_cscColPointers[0] = 0;
        
        
        // ✅ COO 转换
        int* c_cooRowIndices = (int*)malloc(nnz * sizeof(int));
        int* c_cooColIndices = (int*)malloc(nnz * sizeof(int));
        float* c_cooValues = (float*)malloc(nnz * sizeof(float));

        for (int i = 0; i < nnz; i++) {
            c_cooRowIndices[i] = rowIndices[i].intValue;
            c_cooColIndices[i] = colIndices[i].intValue;
            c_cooValues[i] = values[i].floatValue;
        }


        // ✅ HYB 转换
        int ellWidth = 0;
        for (int row = 0; row < numRows; row++) {
            int rowStart = c_cooRowIndices[row];
            int rowEnd = c_cooRowIndices[row + 1];
            ellWidth = MAX(ellWidth, rowEnd - rowStart);
        }
        ellWidth = MIN(ellWidth, 128);

        int* c_hybEllColIndices = (int*)malloc(numRows * ellWidth * sizeof(int));
        float* c_hybEllValues = (float*)malloc(numRows * ellWidth * sizeof(float));
        memset(c_hybEllColIndices, -1, numRows * ellWidth * sizeof(int));

        for (int i = 0; i < nnz; i++) {
            int row = c_cooRowIndices[i];
            for (int j = 0; j < ellWidth; j++) {
                if (c_hybEllColIndices[row * ellWidth + j] == -1) {
                    c_hybEllColIndices[row * ellWidth + j] = c_cooColIndices[i];
                    c_hybEllValues[row * ellWidth + j] = c_cooValues[i];
                    break;
                }
            }
        }

        // 初始化向量 x 和 y
        float* x = (float*)malloc(numCols * sizeof(float));
        float* y = (float*)malloc(numRows * sizeof(float));
        for (int i = 0; i < numCols; i++) x[i] = 1.0f;
        memset(y, 0, numRows * sizeof(float));

        
        
        // ✅ 修正 Metal Heap 可能超出 GPU 限制的问题
        MTLHeapDescriptor *heapDescriptor = [[MTLHeapDescriptor alloc]init];
        size_t totalRequiredSize = (nnz * sizeof(float)) * 4 + (numRows + 1) * sizeof(int) + (numCols + 1) * sizeof(int);
        heapDescriptor.size = MIN(1024 * 1024 * 512, totalRequiredSize);

        // ✅ 允许 Heap 超出时改用普通 MTLBuffer 分配
        bool useHeap = true;
        if (heapDescriptor.size > (1024 * 1024 * 512)) {
            NSLog(@"⚠️ Heap 超出限制，改用普通 MTLBuffer 分配");
            useHeap = false;
        }

        // ✅ 在 Heap 直接分配 GPU 共享缓冲区
       // MTLHeapDescriptor *heapDescriptor = [[MTLHeapDescriptor alloc] init];
       // heapDescriptor.size = 1024 * 1024 * 50; // 50MB
       // heapDescriptor.storageMode = MTLStorageModeShared;
       // id<MTLHeap> heap = [commandQueue.device newHeapWithDescriptor:heapDescriptor];

        // ✅ 在 Heap 直接分配 GPU 共享缓冲区，或者使用普通 MTLBuffer
        // ✅ 在 Heap 直接分配 GPU 共享缓冲区，或使用普通 MTLBuffer
        id<MTLHeap> heap = nil;
        if (useHeap) {
            heap = [device newHeapWithDescriptor:heapDescriptor];
        }
        id<MTLBuffer> csrValuesBuffer = useHeap ?
            [heap newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared];

        id<MTLBuffer> csrRowPointersBuffer = useHeap ?
            [heap newBufferWithLength:(numRows + 1) * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:(numRows + 1) * sizeof(int) options:MTLResourceStorageModeShared];

        id<MTLBuffer> csrColumnIndicesBuffer = useHeap ?
            [heap newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared];

        // ✅ CSC 格式缓冲区
        id<MTLBuffer> cscValuesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared];

        id<MTLBuffer> cscColPointersBuffer = useHeap && heap ?
            [heap newBufferWithLength:(numCols + 1) * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:(numCols + 1) * sizeof(int) options:MTLResourceStorageModeShared];

        id<MTLBuffer> cscRowIndicesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared];
        
        
        // ✅ COO 格式缓冲区
        id<MTLBuffer> cooValuesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared];

        id<MTLBuffer> cooRowIndicesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared];

        id<MTLBuffer> cooColIndicesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared];
        
        
        // ✅ HYB 格式缓冲区（ELL + CSR 混合）
        
        id<MTLBuffer> hybEllValuesBuffer = useHeap && heap ?
            [heap newBufferWithLength:(numRows * ellWidth) * sizeof(float) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:(numRows * ellWidth) * sizeof(float) options:MTLResourceStorageModeShared];

        id<MTLBuffer> hybEllColIndicesBuffer = useHeap && heap ?
            [heap newBufferWithLength:(numRows * ellWidth) * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:(numRows * ellWidth) * sizeof(int) options:MTLResourceStorageModeShared];

        // ✅ HYB 还需要 CSR 部分的存储
        id<MTLBuffer> hybCsrRowPointersBuffer = useHeap && heap ?
            [heap newBufferWithLength:(numRows + 1) * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:(numRows + 1) * sizeof(int) options:MTLResourceStorageModeShared];

        id<MTLBuffer> hybCsrColumnIndicesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(int) options:MTLResourceStorageModeShared];

        id<MTLBuffer> hybCsrValuesBuffer = useHeap && heap ?
            [heap newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared] :
            [device newBufferWithLength:nnz * sizeof(float) options:MTLResourceStorageModeShared];
        
        
        // ✅ 直接 memcpy 避免 `BlitCommand`
       // float *c_values = (float *)malloc(nnz * sizeof(float));
        c_values = (float *)malloc(nnz * sizeof(float)); // ✅ 直接赋值，不重复声明
        memset(c_values, 1, nnz * sizeof(float));
        memcpy(csrValuesBuffer.contents, c_values, nnz * sizeof(float));

        

        // ========================
        // 加载 Metal 内核
        // ========================
        NSError *error = nil;
        NSString *libraryPath = [[NSBundle mainBundle] pathForResource:@"default" ofType:@"metallib"];
        id<MTLLibrary> library = [device newLibraryWithFile:libraryPath error:&error];
        id<MTLFunction> csrKernel = [library newFunctionWithName:@"spmv_csr"];
        id<MTLFunction> cooKernel = [library newFunctionWithName:@"spmv_coo"];
        id<MTLFunction> cscKernel = [library newFunctionWithName:@"spmv_csc"];
        id<MTLFunction> hybKernel = [library newFunctionWithName:@"spmv_hyb"];
        id<MTLComputePipelineState> csrPipeline = [device newComputePipelineStateWithFunction:csrKernel error:&error];
        id<MTLComputePipelineState> cooPipeline = [device newComputePipelineStateWithFunction:cooKernel error:&error];
        id<MTLComputePipelineState> cscPipeline = [device newComputePipelineStateWithFunction:cscKernel error:&error];
        id<MTLComputePipelineState> hybPipeline = [device newComputePipelineStateWithFunction:hybKernel error:&error];

        // ========================
        // 执行计算并记录时间
        // ========================
        double cpuSetupTime, gpuExecutionTime;
        // ✅ 运行 CSR Kernel
        executeKernelWithTiming(@"spmv_csr", @[csrValuesBuffer, csrColumnIndicesBuffer, csrRowPointersBuffer], commandQueue, numRows);
               
               // ✅ 运行 CSC Kernel
        executeKernelWithTiming(@"spmv_csc", @[cscValuesBuffer, cscRowIndicesBuffer, cscColPointersBuffer], commandQueue, numCols);
               
               // ✅ 运行 COO Kernel
        executeKernelWithTiming(@"spmv_coo", @[cooValuesBuffer, cooRowIndicesBuffer, cooColIndicesBuffer], commandQueue, nnz);
               
               // ✅ 运行 HYB Kernel
//        executeKernelWithTiming(@"spmv_hyb", @[csrValuesBuffer, csrColumnIndicesBuffer, csrRowPointersBuffer, cscValuesBuffer, cscRowIndicesBuffer, cscColPointersBuffer], commandQueue, numRows);
        executeKernelWithTiming(@"spmv_hyb", @[hybEllValuesBuffer, hybEllColIndicesBuffer,
                                               hybCsrValuesBuffer, hybCsrColumnIndicesBuffer, hybCsrRowPointersBuffer], commandQueue, numRows);

        // 释放内存
        free(c_values);
        free(c_rowIndices);
        free(c_colIndices);
        free(c_rowPointers);
        free(c_cscColPointers);
        free(c_cscRowIndices);
        free(c_cscValues);
        free(c_hybEllColIndices);
        free(c_hybEllValues);
        free(x);
        free(y);
        
        // ✅ 释放 Metal 资源的方法（ARC 自动管理）
        csrValuesBuffer = nil;
        csrRowPointersBuffer = nil;
        csrColumnIndicesBuffer = nil;

        cscValuesBuffer = nil;
        cscColPointersBuffer = nil;
        cscRowIndicesBuffer = nil;

        cooValuesBuffer = nil;
        cooRowIndicesBuffer = nil;
        cooColIndicesBuffer = nil;

        hybEllValuesBuffer = nil;
        hybEllColIndicesBuffer = nil;
        hybCsrValuesBuffer = nil;
        hybCsrColumnIndicesBuffer = nil;
        hybCsrRowPointersBuffer = nil;
        


        return 0;
    }
}







